import math
import time
from squaternion import Quaternion

import numpy as np
import serial
import torch
import torch.nn as nn
import torch.nn.functional as F

from data_format import Odom, Orientation, Position, Point
from hdware_comm.uart import recv, write
from hdware_comm.lidar import recv_laser_data

# Set the parameters for the implementation
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # cuda or cpu
file_name = "TD3_velodyne"  # name of the file to load the policy from
environment_dim = 20  # laser dimension
robot_dim = 4
state_dim = environment_dim + robot_dim  # state dimension
action_dim = 2  # action dimension
GOAL_REACHED_DIST = 0.3
COLLISION_DIST = 0.5
MAX_V_SPEED = 200  # 最大速度200mm/s
MAX_R_SPEED = 100  # 最大角速度约1.57rad/s
TIME_DELTA = 0.1  # 机器人运动间隔


class Actor(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(Actor, self).__init__()

        self.layer_1 = nn.Linear(state_dim, 800)
        self.layer_2 = nn.Linear(800, 600)
        self.layer_3 = nn.Linear(600, action_dim)
        self.tanh = nn.Tanh()

    def forward(self, s):
        s = F.relu(self.layer_1(s))
        s = F.relu(self.layer_2(s))
        a = self.tanh(self.layer_3(s))
        return a


# TD3 network
class TD3(object):
    def __init__(self, state_dim, action_dim):
        # Initialize the Actor network
        self.actor = Actor(state_dim, action_dim).to(device)

    def get_action(self, state):
        # Function to get the action from the actor
        state = torch.Tensor(state.reshape(1, -1)).to(device)
        return self.actor(state).cpu().data.numpy().flatten()

    def load(self, filename, directory):
        # Function to load network parameters
        self.actor.load_state_dict(
            torch.load("%s/%s_actor.pth" % (directory, filename))
        )


class RealRobotRunner:
    def __init__(self, serial_port='/dev/ttyACM0', baud_rate=115200):
        # 创建并加载TD3网络
        self.network = TD3(state_dim, action_dim)
        try:
            self.network.load(file_name, "./pytorch_models")
            print("Stored model load successfully")
        except:
            raise ValueError("Could not load the stored model")

        # 初始化串口连接
        try:
            self.serial = serial.Serial(serial_port, baud_rate, timeout=5)
            if self.serial.isOpen():
                print("串口连接成功")
            else:
                print("串口连接失败")
        except Exception as e:
            print(f"串口初始化异常: {e}")
            self.serial = None

        self.current_goal = None
        # 位置信息
        self.last_time = time.time()
        self.x = 0.0
        self.y = 0.0
        self.z = 0.0
        self.theta = 0.0
        # 里程计参数
        self.odom = Odom(
            position=Position(0.0, 0.0, 0.0),
            orientation=Orientation(0.0, 0.0, 0.0, 1.0),
        )
        # 雷达参数
        self.velodyne_data = np.ones(environment_dim) * 10  # 初始值均为10，作为最大探测距离
        self.gaps = [[-np.pi / 2 - 0.03, -np.pi / 2 + np.pi / environment_dim]]
        for m in range(environment_dim - 1):
            self.gaps.append(
                [self.gaps[m][1], self.gaps[m][1] + np.pi / environment_dim]
            )
        self.gaps[-1][-1] += 0.03

    def send_action_command(self, action):
        """
        发送动作命令到串口
        返回:
            v_x: mm/s
            v_y: mm/s == 0
            r_z: rad/s
        """
        if self.serial is None or not self.serial.isOpen():
            print("串口未连接，无法发送动作命令")
            return False

        try:
            # 将action转换为适合小车的速度命令
            # action[0]是线速度，范围[0,1]，转换为mm/s
            # action[1]是角速度，范围[-1,1]，转换为rad/s
            v_x = int(action[0] * MAX_V_SPEED)
            v_y = 0  # 假设小车不能横向移动
            r_z = int(action[1] * MAX_R_SPEED)

            write(self.serial, v_x, v_y, r_z)
            return True
        except Exception as e:
            print(f"发送动作命令异常: {e}")
            return False

    def get_odom_data(self):
        """
        从串口获取传感器速度数据
        recv_odom_data:
            数据结构: Odom
        """
        if self.serial is None or not self.serial.isOpen():
            print("串口未连接，无法获取速度数据")
            return None
        try:
            return recv(self.serial)
        except Exception as e:
            print(f"获取速度数据异常: {e}")
            return None

    def get_laser_data(self):
        """
        从lidar.py获取传感器里程计数据
        recv_laser_data:
            数据结构: 点云数据列表，每个点是[x, y, z]坐标
        """
        try:
            return recv_laser_data()
        except Exception as e:
            print(f"获取激光数据异常: {e}")
            return None

    def update_process_odom_data(self, sensor_data):
        """
        用recv接收的速度数据，计算位置和方向信息，更新self.odom
        参数:
            sensor_data: Odom类型
        结果:
            bool: 更新是否成功
        """
        if sensor_data is None:
            return False

        v_x = sensor_data[0]
        r_z = sensor_data[2]

        current_time = time.time()
        dt = current_time - self.last_time
        self.last_time = current_time

        self.theta += r_z * dt
        # 确保theta在[-pi, pi]范围内
        self.theta = math.atan2(math.sin(self.theta), math.cos(self.theta))

        # 更新位置
        self.x += (v_x * math.cos(self.theta) * dt) / 1000.0  # 转换为m
        self.y += (v_x * math.sin(self.theta) * dt) / 1000.0  # 转换为m
        # 创建四元数
        q = euler_to_quaternion(0, 0, self.theta)

        # 更新位置和方向信息
        self.odom.position.x = self.x
        self.odom.position.y = self.y
        self.odom.position.z = self.z

        self.odom.orientation.x = q[0]
        self.odom.orientation.y = q[1]
        self.odom.orientation.z = q[2]
        self.odom.orientation.w = q[3]

        return True

    def update_process_laser_data(self, raw_data):
        """
        处理从recv_laser_data接口获取的原始点云数据，并更新self.velodyne_data
        参数:
            raw_data: 点云数据列表，每个点是[x, y, z]坐标
        结果:
            velodyne_data: 每个角度段中的最小障碍物距离的numpy数组
        """
        # 处理点云中的每个点
        for point in raw_data:
            x, y, z = point

            # 过滤低于特定高度的点
            if z > -0.2:
                # 计算点相对于机器人前进方向的角度
                dot = x * 1 + y * 0
                mag1 = math.sqrt(math.pow(x, 2) + math.pow(y, 2))
                mag2 = math.sqrt(math.pow(1, 2) + math.pow(0, 2))
                beta = math.acos(dot / (mag1 * mag2)) * np.sign(y)

                # 计算到该点的距离
                dist = math.sqrt(x ** 2 + y ** 2 + z ** 2)

                # 将点分配到适当的角度段
                for j in range(len(self.gaps)):
                    if self.gaps[j][0] <= beta < self.gaps[j][1]:
                        self.velodyne_data[j] = min(self.velodyne_data[j], dist)
                        break

    def get_laser_state(self):
        """获取激光雷达状态"""
        v_state = []
        v_state[:] = self.velodyne_data[:]
        laser_state = [v_state]
        return laser_state

    def calculate_theta(self):
        """计算机器人朝向与目标之间的角度"""
        skew_x = self.current_goal.x - self.odom.position.x
        skew_y = self.current_goal.y - self.odom.position.y
        dot = skew_x * 1 + skew_y * 0
        mag1 = math.sqrt(math.pow(skew_x, 2) + math.pow(skew_y, 2))
        mag2 = math.sqrt(math.pow(1, 2) + math.pow(0, 2))
        beta = math.acos(dot / (mag1 * mag2))
        if skew_y < 0:
            if skew_x < 0:
                beta = -beta
            else:
                beta = 0 - beta

        quaternion = Quaternion(
            self.odom.orientation.w,
            self.odom.orientation.x,
            self.odom.orientation.y,
            self.odom.orientation.z,
        )
        euler = quaternion.to_euler(degrees=False)
        angle = round(euler[2], 4)

        theta = beta - angle
        if theta > np.pi:
            theta = np.pi - theta
            theta = -np.pi - theta
        if theta < -np.pi:
            theta = -np.pi - theta
            theta = np.pi - theta
        return theta

    def reset(self):
        """重置环境状态"""
        # 获取传感器数据并更新里程计
        while True:
            sensor_data = self.get_odom_data()
            if sensor_data is not None:
                break
        self.update_process_odom_data(sensor_data)

        # 获取激光雷达状态并更新雷达数据
        while True:
            raw_data = self.get_laser_data()
            if raw_data is not None:
                break
        self.update_process_laser_data(raw_data)
        laser_state = self.get_laser_state()

        # 计算到目标的距离
        distance = np.linalg.norm(
            [self.odom.position.x - self.current_goal.x, self.odom.position.y - self.current_goal.y]
        )

        # 计算朝向与目标之间的角度
        theta = self.calculate_theta()

        # 初始化机器人状态（距离、角度、线速度、角速度）
        robot_state = [distance, theta, 0.0, 0.0]

        # 组合状态
        state = np.append(laser_state, robot_state)
        return state

    @staticmethod
    def observe_collision(laser_data):
        """检测是否发生碰撞"""
        min_laser = min(laser_data)
        if min_laser < COLLISION_DIST:
            return True, True, min_laser
        return False, False, min_laser

    @staticmethod
    def get_reward(target, collision, action, min_laser):
        """计算奖励"""
        if target:
            return 100.0
        elif collision:
            return -100.0
        else:
            r3 = lambda x: 1 - x if x < 1 else 0.0
            return action[0] / 2 - abs(action[1]) / 2 - r3(min_laser) / 2

    def step(self, action):
        """执行一步动作并返回结果"""
        target = False

        # 发送动作命令
        self.send_action_command(action)

        # 等待一小段时间让机器人执行动作
        time.sleep(TIME_DELTA)

        # 获取传感器数据并更新里程计
        while True:
            sensor_data = self.get_odom_data()
            if sensor_data is not None:
                break
        self.update_process_odom_data(sensor_data)

        # 获取激光雷达数据并更新
        while True:
            raw_data = self.get_laser_data()
            if raw_data is not None:
                break
        self.update_process_laser_data(raw_data)
        laser_state = self.get_laser_state()

        # 检测是否发生碰撞
        done, collision, min_laser = self.observe_collision(self.velodyne_data)

        # 计算到目标的距离
        distance = np.linalg.norm(
            [self.odom.position.x - self.current_goal.x, self.odom.position.y - self.current_goal.y]
        )

        # 检测是否到达目标
        if distance < GOAL_REACHED_DIST:
            target = True
            done = True

        # 计算朝向与目标之间的角度
        theta = self.calculate_theta()

        # 更新机器人状态
        robot_state = [distance, theta, action[0], action[1]]

        # 组合状态
        state = np.append(laser_state, robot_state)

        # 计算奖励
        reward = self.get_reward(target, collision, action, min_laser)

        return state, reward, done, target

    def run_episode(self):
        """运行一个episode，直到到达目标点"""
        done = False
        episode_timesteps = 0

        print(f"开始导航至目标点 ({self.current_goal.x}, {self.current_goal.y})")
        state = self.reset()

        while not done:
            # 通过TD3网络选择动作
            action = self.network.get_action(np.array(state))

            # 调整动作范围：线速度[0,1]，角速度[-1,1]
            a_in = [(action[0] + 1) / 2, action[1]]

            # 执行动作并获取新状态
            next_state, reward, done, target = self.step(a_in)
            state = next_state

            if target:
                print("成功到达目标点！")
            elif done and not target:
                print("发生碰撞，任务失败！")

            episode_timesteps += 1

            # 超时检查
            if episode_timesteps > 500:
                print("任务超时！")
                done = True

            # 打印当前状态信息
            distance = np.linalg.norm(
                [self.odom.position.x - self.current_goal.x, self.odom.position.y - self.current_goal.y]
            )
            theta = self.calculate_theta()
            if episode_timesteps % 20 == 0:
                print(f"当前位置: ({self.odom.position.x:.2f}, {self.odom.position.y:.2f}), " +
                      f"距离目标: {distance:.2f}m, 朝向角度: {theta:.2f}rad")

    def run(self):
        """主运行循环"""
        print("输入格式: x y（例如：1.2 3.4）")
        while True:
            try:
                user_input = input("请输入目标点坐标：")
                if user_input.lower() in ["exit", "quit"]:
                    print("已退出运行。")
                    break

                x_str, y_str = user_input.strip().split()
                x, y = float(x_str), float(y_str)

                # 创建目标点并运行导航
                self.current_goal = Point(x=x, y=y)
                self.run_episode()

            except ValueError:
                print("输入格式有误，请输入两个数字，例如：1.0 2.0")
            except KeyboardInterrupt:
                print("\n已手动中断。")
                break
            except Exception as e:
                print(f"发生异常: {e}")
            time.sleep(0.1)

    def close(self):
        """关闭资源"""
        if self.serial and self.serial.isOpen():
            # 停止小车
            try:
                write(self.serial, 0, 0, 0)
            except:
                pass
            self.serial.close()
            print("串口已关闭")


if __name__ == "__main__":
    runner = RealRobotRunner()
    try:
        runner.run()
    finally:
        runner.close()


def euler_to_quaternion(roll, pitch, yaw):
    """
    将欧拉角转换为四元数

    参数:
        roll (float): 绕x轴的旋转角度(弧度)
        pitch (float): 绕y轴的旋转角度(弧度)
        yaw (float): 绕z轴的旋转角度(弧度)

    返回:
        list: 四元数 [x, y, z, w]
    """
    # 计算各个角度的一半的正弦和余弦值
    cy = math.cos(yaw * 0.5)
    sy = math.sin(yaw * 0.5)
    cp = math.cos(pitch * 0.5)
    sp = math.sin(pitch * 0.5)
    cr = math.cos(roll * 0.5)
    sr = math.sin(roll * 0.5)

    # 计算四元数的各个分量
    w = cr * cp * cy + sr * sp * sy
    x = sr * cp * cy - cr * sp * sy
    y = cr * sp * cy + sr * cp * sy
    z = cr * cp * sy - sr * sp * cy

    return [x, y, z, w]
